#! /usr/bin/env python
# coding=utf-8



# 从infernal的结果去除 tRNA和miRNA 而提取 rRNA和snRNA
# 将结果输出为一个gff文件


import os
import shutil
import sys
import argparse


parser = argparse.ArgumentParser(
    description='''从infernal的结果去除 tRNA和miRNA 而提取 rRNA和snRNA
    用法:
    take_rRNA_snRNA_from_infernal.py -i infernal.out -o NIG_ONT_INFERNAL -p NIG_ONT
    由大天才于2021年7月12日创建于浙江农业大学''')

parser.add_argument('-i',
                help='必须给定，输入infernal的结果')


parser.add_argument('-o',
                help='必须给定，输出文件的路径')

parser.add_argument('-p',
                help='前缀，默认为ncRNA')


args = parser.parse_args()

if not args.i or not args.o:
    parser.print_help()
    sys.exit()


infile = args.i

outfile  = args.o



if not args.p:
    prefix = 'ncRNA'
else:
    prefix = args.p




import re

# infile = r'infernal.out'


# outfile = r'out'




remove_lista = ["miR","tRNA"]




snoRNA_info = ['snoRNA','small nucleolar']

snoRNA_lista = []


rRNA_info = ['rRNA','ribosomal RNA']

rRNA_lista = []

snRNA_info = ['snRNA','spliceosomal RNA','7SK RNA','Wolbachia sRNA']

snRNA_lista = []

with open(infile) as fila:
	for i in fila:

		i = i.strip().split('#')[0]

		k = re.split(r'\s+',i.strip())

		if len(k)>5:
			
			des = re.split(r'\s\s+',i.strip())[-1]


			t = 0
			for j in remove_lista:
				if des.upper().find(j.upper())!=-1:
					t = 1
					break
			if t==1:
				continue

			# 寻找 snoRNA

			t = 0
			for j in snoRNA_info:
				if des.upper().find(j.upper())!=-1:
					t = 1
					break
			if t ==1:
				snoRNA_lista.append(k+[des])
				continue
			
		
			# 寻找 rRNA

			t = 0
			for j in rRNA_info:
				if des.upper().find(j.upper())!=-1:
					t = 1
					break
			if t ==1:
				rRNA_lista.append(k+[des])
				continue



			t = 0
			for j in snRNA_info:
				if des.upper().find(j.upper())!=-1:
					t = 1
					break
			if t ==1:
				snRNA_lista.append(k+[des])
				continue

n = 1


with open(outfile+'_snoRNA.gff' ,'w') as fita:
	fita.write('##gff-version 3.2.1\n')
	for i in snoRNA_lista:
		scaf = i[2]
		begin = min([int(i[7]),int(i[8])])
		end = max([int(i[7]),int(i[8])])
		d = i[8]

		name = prefix + '_snoRNA_'+ str(n) 

		n += 1

		fita.write(scaf+'\t'+'infernal\tsnoRNA'+'\t'+str(begin)+'\t'+str(end)+'\t.\t'+d+'\t.\tID='+name+'\n')


with open(outfile+'_rRNA.gff' ,'w') as fita:
	fita.write('##gff-version 3.2.1\n')
	for i in rRNA_lista:
		scaf = i[2]
		begin = min([int(i[7]),int(i[8])])
		end = max([int(i[7]),int(i[8])])
		d = i[8]

		name = prefix + '_rRNA_'+ str(n) 

		n += 1

		fita.write(scaf+'\t'+'infernal\trRNA'+'\t'+str(begin)+'\t'+str(end)+'\t.\t'+d+'\t.\tID='+name+'\n')

with open(outfile+'_snRNA.gff' ,'w') as fita:
	fita.write('##gff-version 3.2.1\n')
	for i in snRNA_lista:
		scaf = i[2]
		begin = min([int(i[7]),int(i[8])])
		end = max([int(i[7]),int(i[8])])
		d = i[8]

		name = prefix + '_snRNA_'+ str(n) 

		n += 1

		fita.write(scaf+'\t'+'infernal\tsnRNA'+'\t'+str(begin)+'\t'+str(end)+'\t.\t'+d+'\t.\tID='+name+'\n')